"""
主程序入口
基于Transformer的复数信号稀疏表示学习系统
"""
import argparse
import os
from config import Config
from utils import set_seed, print_model_info
from trainer import Trainer

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='Transformer自编码器训练程序')
    
    parser.add_argument('--config', type=str, default=None,
                       help='配置文件路径')
    parser.add_argument('--train', action='store_true',
                       help='训练模型')
    parser.add_argument('--test', action='store_true',
                       help='测试模型')
    parser.add_argument('--evaluate', action='store_true',
                       help='评估序列数据')
    parser.add_argument('--load_model', type=str, default=None,
                       help='加载预训练模型路径')
    parser.add_argument('--save_model', type=str, default=None,
                       help='保存模型路径')
    parser.add_argument('--epochs', type=int, default=None,
                       help='训练轮数')
    parser.add_argument('--batch_size', type=int, default=4,
                       help='批次大小')
    parser.add_argument('--learning_rate', type=float, default=None,
                       help='学习率')
    parser.add_argument('--force_retrain', action='store_true',
                       help='强制重新训练')
    parser.add_argument('--eval_start_idx', type=int, default=1500,
                       help='评估序列的起始索引')
    
    return parser.parse_args()

def update_config_from_args(config, args):
    """根据命令行参数更新配置"""
    if args.epochs is not None:
        config.EPOCHS = args.epochs
    if args.batch_size is not None:
        config.BATCH_SIZE = args.batch_size
    if args.learning_rate is not None:
        config.LEARNING_RATE = args.learning_rate
    if args.force_retrain:
        config.FORCE_RETRAIN = True
    
    return config

def main():
    """主函数"""
    # 解析命令行参数
    args = parse_arguments()
    
    # 创建配置
    config = Config()
    config = update_config_from_args(config, args)
    
    # 设置随机种子
    set_seed(config.SEED)
    
    # 创建必要的目录
    os.makedirs(os.path.dirname(config.MODEL_SAVE_PATH), exist_ok=True)
    
    print("=" * 60)
    print("基于Transformer的复数信号稀疏表示学习系统")
    print("=" * 60)
    print(f"设备: {config.DEVICE}")
    print(f"随机种子: {config.SEED}")
    print(f"批次大小: {config.BATCH_SIZE}")
    print(f"序列长度: {config.SEQ_LEN}")
    print(f"端元个数: {config.K}")
    print(f"稀疏性约束: {config.K_NONZERO}")
    print("=" * 60)
    
    # 创建训练器
    trainer = Trainer(config)
    
    # 打印模型信息
    print_model_info(trainer.get_model())
    print("=" * 60)
    
    # 如果指定了加载模型路径
    if args.load_model:
        trainer.load_model(args.load_model)
    
    # 执行相应的操作
    if args.train or (not args.test and not args.evaluate):
        # 默认执行训练
        print("开始训练...")
        loss_values = trainer.train()
        
        if args.save_model:
            trainer.save_model(args.save_model)
    
    if args.test or (not args.train and not args.evaluate):
        # 如果只指定了test或者没有指定任何操作，执行测试
        print("\n开始测试...")
        rmse_list, test_loss = trainer.test()
    
    if args.evaluate:
        # 评估序列数据
        print(f"\n开始评估序列数据...")
        Y_reconstructed, X, mse_total, rmse = trainer.evaluate_sequence(args.eval_start_idx)
    
    print("\n程序执行完成！")

if __name__ == "__main__":
    main() 